#include "geometry_msgs/msg/pose_array.hpp"
#include "rclcpp/rclcpp.hpp"
#include "std_msgs/msg/u_int8_multi_array.hpp"
#include <opencv2/opencv.hpp>

#include <fstream>

#include "final_exam/armor_detect.h"
#include "final_exam/json.hpp"
#include "final_exam/param.h"

nlohmann::json param::settings;

class CenterNode : public rclcpp::Node {
public:
  CenterNode(std::string name) : Node(name) {
    RCLCPP_INFO(this->get_logger(), "Center Node '%s' has been started.",
                name.c_str());

    // 订阅图像话题（消息为 UInt8MultiArray，包含 JPEG 字节流）
    image_subscription_ =
        this->create_subscription<std_msgs::msg::UInt8MultiArray>(
            "video_frames", 10,
            [&](const std_msgs::msg::UInt8MultiArray::SharedPtr msg) {
              imageCallback(msg);
            });
    // 发布标注后图像话题（改用 UInt8MultiArray 传输二进制数据）
    annotated_image_publisher_ =
        this->create_publisher<std_msgs::msg::UInt8MultiArray>(
            "annotated_image_topic", 10);
    // 发布装甲板位置数据话题
    armor_pose_publisher_ =
        this->create_publisher<geometry_msgs::msg::PoseArray>(
            "armor_pose_topic", 10);
    // 设置日志级别
    this->get_logger().set_level(rclcpp::Logger::Level::Info);
  }

  // 加载参数文件
  // TODO: 修改为 ROS2 参数读取
  void loadArmorDetectParams() {
    // 读取相机内参
    std::vector<std::vector<double>> camera_matrix_vec =
        param::settings["camera_matrix"];
    cv::Mat camera_matrix(3, 3, CV_64F);
    for (int i = 0; i < 3; i++) {
      for (int j = 0; j < 3; j++) {
        camera_matrix.at<double>(i, j) = camera_matrix_vec[i][j];
      }
    }
    std::vector<double> dist_coeffs_vec = param::settings["dist_coeffs"];
    cv::Mat dist_coeffs(dist_coeffs_vec);
    armor_detector_.setCameraParams(camera_matrix, dist_coeffs);

    // 读取物体点
    std::vector<std::vector<double>> object_points_vec =
        param::settings["object_points"];
    std::vector<cv::Point3f> object_points;
    for (const auto &point : object_points_vec) {
      if (point.size() == 3) {
        object_points.emplace_back(point[0], point[1], point[2]);
      }
    }
    armor_detector_.setObjectPoints(object_points);

    // 设置敌人颜色
    std::string enemy_color = param::settings["enemy"];
    if (enemy_color == "RED") {
      armor_detector_.setEnemy(param::RED);
    } else if (enemy_color == "BLUE") {
      armor_detector_.setEnemy(param::BLUE);
    } else {
      RCLCPP_WARN(this->get_logger(),
                  "Unknown enemy color '%s', defaulting to RED.",
                  enemy_color.c_str());
      armor_detector_.setEnemy(param::RED);
    }

    RCLCPP_INFO(this->get_logger(), "ArmorDetect parameters loaded.");
  }

private:
  // 装甲板检测器
  ArmorDetect armor_detector_;
  // 图像话题订阅者
  rclcpp::Subscription<std_msgs::msg::UInt8MultiArray>::SharedPtr
      image_subscription_;
  // 标注后图像话题发布者
  rclcpp::Publisher<std_msgs::msg::UInt8MultiArray>::SharedPtr
      annotated_image_publisher_;
  // 装甲板位置数据话题发布者
  rclcpp::Publisher<geometry_msgs::msg::PoseArray>::SharedPtr
      armor_pose_publisher_;

  // 话题回调函数
  void imageCallback(const std_msgs::msg::UInt8MultiArray::SharedPtr msg) {
    if (!msg) {
      RCLCPP_WARN(this->get_logger(), "Received null image message");
      return;
    }

    if (msg->data.empty()) {
      RCLCPP_WARN(this->get_logger(), "Received empty image data");
      return;
    }

    // 直接使用 UInt8MultiArray 的 data 字段作为 uchar 向量
    const auto &buf = msg->data;

    // 使用 OpenCV 解码 JPEG 数据；解码后的图像为 BGR
    cv::Mat img = cv::imdecode(buf, cv::IMREAD_COLOR);
    RCLCPP_DEBUG(this->get_logger(),
                 "Decoding JPEG image from message of size %zu bytes",
                 buf.size());
    if (img.empty()) {
      RCLCPP_WARN(this->get_logger(),
                  "Failed to decode JPEG image from message");
      return;
    }

    // 成功解码
    RCLCPP_DEBUG(this->get_logger(),
                 "Decoded image: width=%d height=%d channels=%d", img.cols,
                 img.rows, img.channels());
    // 装甲板检测
    cv::Mat dst;
    armor_detector_.detect(img, dst);

    // 发布标注后图像
    std_msgs::msg::UInt8MultiArray annotated_msg;
    std::vector<uchar> encoded_buf;
    if (!cv::imencode(".jpg", dst, encoded_buf)) {
      RCLCPP_WARN(this->get_logger(),
                  "Failed to encode annotated image to JPEG");
    } else {
      annotated_msg.data.assign(encoded_buf.begin(), encoded_buf.end());
      annotated_image_publisher_->publish(annotated_msg);
    }

    // 发布装甲板位置数据
    armor_pose_publisher_->publish(armor_detector_.getArmorPosesMsg());
  }
};

void loadJson() {
  // 加载参数文件
  std::ifstream param_file("settings.json");
  if (!param_file.is_open()) {
    throw std::runtime_error("Failed to open parameter file.");
  }

  param_file >> param::settings;
  param_file.close();
}

int main(int argc, char **argv) {
  // 初始化 ROS2
  rclcpp::init(argc, argv);

  // 创建 CenterNode 节点
  auto center_node = std::make_shared<CenterNode>("center_node");

  // 加载参数文件
  loadJson();
  center_node->loadArmorDetectParams();

  // 运行节点
  rclcpp::spin(center_node);

  // 关闭 ROS2
  rclcpp::shutdown();
  return 0;
}